// Copyright 2017 guangbo. All rights reserved.

//rpc编码解码模块
package lib_rpc

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"github.com/golang/protobuf/proto"
	"io"
	"sync"
)

func writeFrame(w *bufio.Writer, buf []byte) error {
	var lenBuf [2]byte
	binary.BigEndian.PutUint16(lenBuf[:], uint16(len(buf)))
	_, err := w.Write(lenBuf[:])
	if err != nil {
		return err
	}
	_, err = w.Write(buf)
	return err
}

func readFrame(r *bufio.Reader) ([]byte, error) {
	buff := make([]byte, 2)
	_, err := io.ReadFull(r, buff)
	if err != nil {
		return nil, err
	}
	len := binary.BigEndian.Uint16(buff)

	buff = make([]byte, len)
	_, err = io.ReadFull(r, buff)
	if err != nil {
		return nil, err
	}

	return buff, nil
}

func encode(w *bufio.Writer, m interface{}) error {
	if pb, ok := m.(proto.Message); ok {
		buf, err := proto.Marshal(pb)
		if err != nil {
			return err
		}
		return writeFrame(w, buf)
	}
	return fmt.Errorf("%T does not implement proto.Message", m)
}

func decode(r *bufio.Reader, m interface{}) error {
	buff, err := readFrame(r)
	if err != nil {
		return err
	}

	return proto.Unmarshal(buff, m.(proto.Message))
}

type PbClientCodec struct {
	req  ReqHeader
	resp RspHeader
	c    io.ReadWriteCloser
	w    *bufio.Writer
	r    *bufio.Reader
}

func NewPbClientCodec(rwc io.ReadWriteCloser) ClientCodec {
	return &PbClientCodec{
		r: bufio.NewReaderSize(rwc, 4096),
		w: bufio.NewWriterSize(rwc, 4096),
		c: rwc,
	}
}

func (c *PbClientCodec) WriteRequest(r *Request, body interface{}) error {
	c.req.Reset()
	c.req.Method = r.ServiceMethod
	c.req.Seq = r.Seq
	err := encode(c.w, &c.req)
	if err != nil {
		return err
	}
	if err = encode(c.w, body); err != nil {
		return err
	}
	return c.w.Flush()
}

func (c *PbClientCodec) ReadResponseHeader(r *Response) error {
	c.resp.Reset()
	err := decode(c.r, &c.resp)
	if err != nil {
		return err
	}
	r.ServiceMethod = c.resp.Method
	r.Seq = c.resp.Seq
	r.Error = c.resp.Error
	r.Ret = c.resp.Ret
	return nil
}

func (c *PbClientCodec) ReadResponseBody(body interface{}) error {
	return decode(c.r, body)
}

func (c *PbClientCodec) WriteByteRequest(r *Request, buf []byte) error {
	c.req.Reset()
	c.req.Method = r.ServiceMethod
	c.req.Seq = r.Seq
	err := encode(c.w, &c.req)
	if err != nil {
		return err
	}
	if err = writeFrame(c.w, buf); err != nil {
		return err
	}
	return c.w.Flush()
}

func (c *PbClientCodec) ReadByteResponseBody() ([]byte, error) {
	return readFrame(c.r)
}

func (c *PbClientCodec) Close() error {
	return c.c.Close()
}

type PbServerCodec struct {
	mu   sync.Mutex // exclusive writer lock
	req  ReqHeader
	resp RspHeader
	w    *bufio.Writer
	r    *bufio.Reader

	c io.Closer
}

func NewPbServerCodec(rwc io.ReadWriteCloser) ServerCodec {
	return &PbServerCodec{
		r: bufio.NewReaderSize(rwc, 4096),
		w: bufio.NewWriterSize(rwc, 4096),
		c: rwc,
	}
}

func (c *PbServerCodec) WriteResponse(resp *Response, body interface{}) error {
	c.mu.Lock()
	c.resp.Method = resp.ServiceMethod
	c.resp.Seq = resp.Seq
	c.resp.Error = resp.Error
	c.resp.Ret = resp.Ret

	err := encode(c.w, &c.resp)
	if err != nil {
		c.mu.Unlock()
		return err
	}
	if err = encode(c.w, body); err != nil {
		c.mu.Unlock()
		return err
	}
	err = c.w.Flush()
	c.mu.Unlock()
	return err
}

func (c *PbServerCodec) ReadRequestHeader(req *Request) error {
	c.req.Reset()

	err := decode(c.r, &c.req)
	if err != nil {
		return err
	}
	req.ServiceMethod = c.req.Method
	req.Seq = c.req.Seq
	return nil
}

func (c *PbServerCodec) ReadRequestBody(body interface{}) error {
	return decode(c.r, body)
}

func (c *PbServerCodec) Close() error { return c.c.Close() }
